# 可视化原图和篡改位置
import os
import cv2
import six
import lmdb
import argparse
from PIL import Image
import numpy as np

parser = argparse.ArgumentParser()
parser.add_argument('--input', type=str,
                    default='DocTamperV1-FCD')  ## path to the dataset dir, which contains data.mdb and lock.mdb
parser.add_argument('--i', type=int, default=0)  # index of the image to be visualized
parser.add_argument('--output', type=str,
                    default='/kaggle/working/doc-tamper2/pic')  # index of the image to be visualized
args = parser.parse_args()
a = lmdb.open(args.input, readonly=True, lock=False, readahead=False, meminit=False)
index = args.i
with a.begin(write=False) as txn:
    img_key = 'image-%09d' % index
    imgbuf = txn.get(img_key.encode('utf-8'))
    buf = six.BytesIO()
    buf.write(imgbuf)
    buf.seek(0)
    im = Image.open(buf)
    print(str(args.output))
    im.save(args.output + f'a{index}.jpg')
    lbl_key = 'label-%09d' % index
    lblbuf = txn.get(lbl_key.encode('utf-8'))
    mask = cv2.imdecode(np.frombuffer(lblbuf, dtype=np.uint8), 0)
    print(mask.max())
    if mask.max() == 1:
        mask = mask * 255
    cv2.imwrite(args.output + f'a{index}.png', mask)

if __name__ == '__main__':
    s = 'python vizlmdb.py --output /kaggle/working/doc-tamper2/pic --input /kaggle/input/mytext/DocTamperV1-FCD --i '
    list = []
    for i in range(500):
        buf = s + str(i) + '\n'
        list.append(buf)
    with open('list.sh', 'w') as f:
        for z in list:
            f.write(z)
